# -*- coding: utf-8 -*-

import numpy as np
import ptycho

# General parameters
verbose_level = 3
data_type = 'single'
run_label = 'linear'

# Physical dimensions
energy = 6.2                # Energy (in keV)
z = 7.046                  # Distance from object to screen 
ds = 172e-6                 # Camera pixel size

scans = ['simul_lownoise_sp_displacements']
asize = np.array((256, 256))
ctr = (129, 129)
upsample_data = None
flip_data = False

# Scan parameters
scan_type = 'round_roi'                     # ['round', 'raster', 'round_roi']
radius_in = 0                           # interior radius of the round scan
radius_out = 1.5e-6                       # exterior radius of the round scan
nr = 6                                 # number of intervals (# of shells - 1)
nth = 5                                 # number of points in the first shell
lx = 12e-6                              # Width of ROI for round_ROI scans
ly = 12e-6                              # Height or ROI for round_ROI scans
dr = 0.7e-6                             # Shell spacing for round_ROI scans
nx = 10                                 # raster scan: number of steps in x
ny = 10                                 # raster scan: number of steps in y
step_size_x = 1e-6                      # raster scan: step size (grid spacing)
step_size_y = 1e-6                      # raster scan: step size (grid spacing)

# Multiple-scan pattern.
scan_pattern = (0, 0)   # Unique probe and object over all scans
#scan_pattern = (1,0)  # one probe per scan, one (shared) object
#scan_pattern = (0,1)  # single (shared) probe, multiple objects 
#scan_pattern = (3,0)  # single probe over three consecutive scans, single object
#scan_pattern = None    # For any more complicated case
#scan_pattern_probe_indices = [...]
#scan_pattern_object_indices = [...]


# Initial probe
initial_probe_type = 'file'
initial_probe_file = './probe_0.h5'
probe_diameter = 1.5e-6
probe_propdist = 0.7e-3

# Initial object
initial_object_type = 'ones'
#initial_object_file = '../sample/simul_weak_solution.h5'

#average_probes = True
#average_probe_amp = 1e-2
dp_shift = False

subpix = False
subpix_method = 'linear'
subpix_start = 10

subpix_disp = True
subpix_disp_start = 20

probe_support = False
probe_support_start = 0
probe_support_area = .6                    # Area (relative to total area) of probe support
probe_before_object = False

probe_antialiasing = None

# I/O, interactions
pathdir_patt = './%(scan)s/'
datafile_patt = '%(path)s%(scan)s_data_%(a0)03dx%(a1)03d.h5'  # datafile pattern
save_dir_patt = './'
dump_interval = 1000                 # Interval for dumping intermediate results
dump_patt = '%(run_name)s_dump.hd5'
plot_interval = 10                    # interval for plotting
doplot = False
last_plot = False                   # Whether to show a (blocking) last plot of the reconstruction
dump_plot = True                    # Dumping an image file of the plot
dump_plot_patt = dump_plot_patt = scans[0] + '_dump_DM_fourier_sub_%d.pdf'

# Reconstruction (general)
# Not used for now
# algorithm = 'ML'
# algorithm = 'DM'
numit = 200
numit_DM = 40
numit_ML = 50

# Reconstruction (DM)
fourier_relax_factor = 0.05
pbound = None                      # Power bound to be applied in the Fourier projection
clip_object = False
clip_max = 1.0
clip_min = 0.01
average_start = 100
average_interval = 10
probe_change_start = 2
DM_smooth_amplitude = 0.1

# Reconstruction (ML)
quad_interval = 0                    # Interval to compute real gradient instead of quadratic approximation.
scale_precond = True
ML_type = 'Gauss'
reg_del2 = True
reg_del2_amplitude = 0.01

###############################################################
###############################################################
# Prepare parameters

p = ptycho.prepare_params(globals())
ptycho.verbose(1, ptycho.print_summary(p))
# Run reconstruction
p = ptycho.ptycho_DM(p, numit=numit_DM, subpix=False)
p = ptycho.ptycho_ML(p, numit=numit_ML, subpix=True, dump_plot_patt=scans[0] + '_dump_ML_fourier_sub_%d.pdf')
